import tensorflow as tf
import numpy as np
from PIL import Image
import logging
from typing import List

logging.basicConfig(level=logging.INFO)

class TensorflowRecognizer:
    def __init__(self, model_path: str = 'captcha_recognition/models/captcha_model.h5'):
        """
        初始化TensorFlow验证码识别器
        Args:
            model_path: 模型文件路径
        """
        self.image_height = 64
        self.image_width = 128
        self.characters = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
        
        try:
            logging.info(f"加载模型: {model_path}")
            self.model = tf.keras.models.load_model(model_path)
        except Exception as e:
            logging.error(f"加载模型失败: {str(e)}")
            raise
        
    def _preprocess_image(self, image_path: str) -> np.ndarray:
        """
        预处理图片
        Args:
            image_path: 图片路径
        Returns:
            np.ndarray: 预处理后的图片数组
        """
        try:
            image = Image.open(image_path)
            image = image.resize((self.image_width, self.image_height))
            image = np.array(image) / 255.0
            return np.expand_dims(image, axis=0)
        except Exception as e:
            logging.error(f"图片预处理失败: {str(e)}")
            raise
            
    def _decode_predictions(self, predictions: List[np.ndarray]) -> str:
        """
        解码模型预测结果
        Args:
            predictions: 模型预测输出
        Returns:
            str: 识别出的验证码文本
        """
        result = ''
        for pred in predictions:
            char_idx = np.argmax(pred[0])
            if char_idx < len(self.characters):
                result += self.characters[char_idx]
        return result
        
    async def recognize(self, image_path: str) -> str:
        """
        识别验证码
        Args:
            image_path: 验证码图片路径
        Returns:
            str: 识别出的验证码文本
        """
        try:
            logging.info(f"开始识别验证码: {image_path}")
            
            # 预处理图片
            image = self._preprocess_image(image_path)
            
            # 预测
            predictions = self.model.predict(image)
            
            # 解码预测结果
            result = self._decode_predictions(predictions)
            
            logging.info(f"验证码识别结果: {result}")
            return result
            
        except Exception as e:
            logging.error(f"验证码识别失败: {str(e)}")
            raise